cannot iterate over a tensor with unknown first dimension